import json
import os
import time

from util.llm import http_request

MODEL_NAME = os.getenv('MODEL_NAME', "deepseek-chat")

SYSTEM_PROMPT = """你是一名演练脚本生成专家，需要根据设备、场景描述、应急预案和相关参考文档来生成详细的演练脚本。

工作流：
1. 理解用户提供的设备、场景描述和应急预案
2. 分析应急预案中的各个字段内容
3. 参考相关的演练脚本文档
4. 为每个脚本节点生成具体的内容
5. 对话节点要包含具体的对话内容
6. 操作节点要包含具体的操作步骤
7. 内容要与应急预案保持一致
8. 必须使用"scriptNode"作为根字段名，不要使用其他字段名

输入信息：
设备：{}
场景描述：{}
应急预案：{}

相关参考文档：
{}

请根据以上信息生成演练脚本，确保内容准确、详细且符合实际演练需求。"""

USER_PROMPT = """请根据设备'{equipment}'和场景描述'{scene_desc}'，以及提供的应急预案和相关参考文档，为演练脚本节点生成具体的内容。

要求：
1. 脚本内容要符合实际演练场景
2. 对话节点的内容应该是具体的对话内容
3. 操作节点的内容应该是具体的操作步骤
4. 内容要与应急预案的相应字段对应
5. 参考相关文档中的最佳实践
6. 必须使用"scriptNode"作为根字段名，不要使用"应急演练脚本"或其他字段名
"""


async def generate_drill_script(equipment: str, scene_desc: str, plan: dict, script_node: list, related_docs: list = None):
    """
    生成演练脚本
    
    Args:
        equipment: 设备
        scene_desc: 场景描述
        plan: 应急预案
        script_node: 原始脚本节点结构
        related_docs: 相关参考文档列表
    
    Returns:
        生成的演练脚本
    """
    
    # 记录总体开始时间
    total_start_time = time.time()
    
    print(f"🔍 generate_drill_script开始处理")
    print(f"🔍 设备: {equipment}")
    print(f"🔍 场景描述: {scene_desc}")
    print(f"🔍 预案字段数: {len(plan)}")
    print(f"🔍 脚本节点数: {len(script_node)}")
    print(f"🔍 相关文档数: {len(related_docs) if related_docs else 0}")
    
    try:
        # 记录提示词构建开始时间
        prompt_start_time = time.time()
        
        # 构建系统提示词
        plan_text = json.dumps(plan, ensure_ascii=False, indent=2)
        
        # 构建相关文档文本
        if related_docs and len(related_docs) > 0:
            docs_text = "参考文档：\n"
            for i, doc in enumerate(related_docs, 1):
                docs_text += f"{i}. {doc.name}\n{doc.content[:500]}...\n\n"
        else:
            docs_text = "无相关参考文档"
        
        system_prompt = SYSTEM_PROMPT.format(equipment, scene_desc, plan_text, docs_text)
        
        # 构建用户查询
        query = USER_PROMPT.format(equipment=equipment, scene_desc=scene_desc)
        
        prompt_end_time = time.time()
        prompt_duration = prompt_end_time - prompt_start_time
        
        print(f"🔍 系统提示词长度: {len(system_prompt)}")
        print(f"🔍 用户查询: {query}")
        print(f"⏱️ 提示词构建耗时: {prompt_duration:.3f}秒")
        
        # 记录LLM调用开始时间
        llm_start_time = time.time()
        
        # 调用LLM生成内容
        result = await http_request(
            system_prompt,
            history=[],
            prompt=query,
            model=MODEL_NAME
        )
        
        llm_end_time = time.time()
        llm_duration = llm_end_time - llm_start_time
        
        print(f"✅ 演练脚本生成成功，结果长度: {len(result)}")
        print(f"🔍 原始结果预览: {result[:200]}...")
        print(f"⏱️ LLM调用耗时: {llm_duration:.3f}秒")
        
    except Exception as e:
        print(f"❌ 演练脚本生成失败: {e}")
        import traceback
        traceback.print_exc()
        raise e
    
    # 记录JSON解析开始时间
    json_start_time = time.time()
    
    # 尝试解析JSON结果
    print(f"🔍 开始解析JSON结果")
    try:
        # 移除可能的markdown代码块标记
        result = result.strip()
        print(f"🔍 去除首尾空白后的结果长度: {len(result)}")
        
        if result.startswith("```json"):
            result = result[7:]
            print(f"🔍 移除```json标记")
        elif result.startswith("```"):
            result = result[3:]
            print(f"🔍 移除```标记")
        if result.endswith("```"):
            result = result[:-3]
            print(f"🔍 移除结尾```标记")
        
        print(f"🔍 清理后的结果长度: {len(result)}")
        print(f"🔍 清理后的结果预览: {result[:200]}...")
        
        # 解析JSON
        script_data = json.loads(result)
        print(f"✅ JSON解析成功")
        print(f"🔍 解析后的数据结构: {list(script_data.keys())}")
        
        # 验证返回结构 - 处理可能的字段名变化
        if 'scriptNode' in script_data:
            print(f"✅ 找到scriptNode字段")
        elif '应急演练脚本' in script_data:
            print(f"⚠️ 找到'应急演练脚本'字段，转换为'scriptNode'")
            script_data['scriptNode'] = script_data.pop('应急演练脚本')
        elif '演练脚本' in script_data:
            print(f"⚠️ 找到'演练脚本'字段，转换为'scriptNode'")
            script_data['scriptNode'] = script_data.pop('演练脚本')
        else:
            print(f"❌ 返回结果缺少scriptNode字段，可用字段: {list(script_data.keys())}")
            # 尝试修复：如果只有一个字段，可能是scriptNode
            if len(script_data) == 1:
                key = list(script_data.keys())[0]
                print(f"⚠️ 尝试将唯一字段'{key}'作为scriptNode")
                script_data['scriptNode'] = script_data.pop(key)
            else:
                raise Exception(f"返回结果缺少scriptNode字段，可用字段: {list(script_data.keys())}")
        
        # 验证scriptNode是否为列表
        if not isinstance(script_data['scriptNode'], list):
            print(f"❌ scriptNode不是列表格式: {type(script_data['scriptNode'])}")
            raise Exception("scriptNode必须是列表格式")
        
        # 验证每个节点是否包含必要字段
        for i, node in enumerate(script_data['scriptNode']):
            if not isinstance(node, dict):
                print(f"❌ 节点{i}不是字典格式")
                raise Exception(f"节点{i}必须是字典格式")
            
            required_fields = ['nodeOrder', 'nodeName', 'nodeType', 'children']
            missing_fields = [field for field in required_fields if field not in node]
            if missing_fields:
                print(f"❌ 节点{i}缺少必要字段: {missing_fields}")
                raise Exception(f"节点{i}缺少必要字段: {missing_fields}")
            
            # 验证children字段
            if not isinstance(node['children'], list):
                print(f"❌ 节点{i}的children不是列表格式")
                raise Exception(f"节点{i}的children必须是列表格式")
        
        # 记录JSON解析结束时间
        json_end_time = time.time()
        json_duration = json_end_time - json_start_time
        
        # 记录总体结束时间
        total_end_time = time.time()
        total_duration = total_end_time - total_start_time
        
        print(f"✅ 演练脚本生成完成，节点数量: {len(script_data['scriptNode'])}")
        print(f"⏱️ JSON解析耗时: {json_duration:.3f}秒")
        print(f"⏱️ 总体耗时: {total_duration:.3f}秒")
        print(f"📊 耗时分析:")
        print(f"   - 提示词构建: {prompt_duration:.3f}秒 ({(prompt_duration/total_duration)*100:.1f}%)")
        print(f"   - LLM调用: {llm_duration:.3f}秒 ({(llm_duration/total_duration)*100:.1f}%)")
        print(f"   - JSON解析: {json_duration:.3f}秒 ({(json_duration/total_duration)*100:.1f}%)")
        
        return script_data
        
    except json.JSONDecodeError as e:
        print(f"❌ JSON解析错误: {e}")
        print(f"❌ 原始结果: {result}")
        
        # 如果JSON解析失败，返回错误信息
        raise Exception(f"演练脚本生成失败，JSON解析错误: {e}")
    except Exception as e:
        print(f"❌ 处理失败: {e}")
        raise e


async def generate_drill_script_stream(equipment: str, scene_desc: str, plan: dict, script_node: list):
    """
    流式生成演练脚本
    
    Args:
        equipment: 设备
        scene_desc: 场景描述
        plan: 应急预案
        script_node: 原始脚本节点结构
    
    Yields:
        生成的演练脚本内容片段
    """
    
    print(f"🔍 generate_drill_script_stream开始处理")
    print(f"🔍 设备: {equipment}")
    print(f"🔍 场景描述: {scene_desc}")
    print(f"🔍 预案字段数: {len(plan)}")
    print(f"🔍 脚本节点数: {len(script_node)}")
    
    try:
        # 构建系统提示词
        plan_text = json.dumps(plan, ensure_ascii=False, indent=2)
        system_prompt = SYSTEM_PROMPT.format(equipment, scene_desc, plan_text, "无相关参考文档") # Placeholder for related docs
        
        print(f"🔍 系统提示词长度: {len(system_prompt)}")
        
        # 构建用户查询
        query = USER_PROMPT.format(equipment=equipment, scene_desc=scene_desc)
        
        print(f"🔍 用户查询: {query}")
        
        # 流式调用LLM生成内容
        async for chunk in http_request( # Changed from llm.sse_request to http_request
            system_prompt,
            history=[],
            prompt=query,
            model=MODEL_NAME
        ):
            yield chunk
            
    except Exception as e:
        print(f"❌ 流式演练脚本生成失败: {e}")
        import traceback
        traceback.print_exc()
        yield f"生成失败: {e}" 